﻿using FireChat.Share.Web.Package;
using Newtonsoft.Json;
using System.Net.WebSockets;
using System.Text;

namespace FireChat.Share.Web.Util;


/// <summary>
/// WebSocket 脚手架
/// </summary>
public static class WebSocketScaffold
{
    /// <summary>
    /// TODO: 消息发送锁 不是很稳定
    /// </summary>
    private static readonly Dictionary<WebSocket, SemaphoreSlim> _sendLock = new();

    /// <summary>
    /// 单个接收
    /// </summary>
    private static readonly Dictionary<WebSocket, bool> _singleReceive = new();



    /// <summary>
    /// 接收数据包 该方法在 (<see cref="WebSocket.State"/>==<see cref="WebSocketState.Closed"/>) 停止之前不会退出
    /// </summary>
    public static Task ReceivePackageAsync<T>(
        this WebSocket webSocket, Action<T> receivedCallback, Func<Exception, bool> excetionProcess) where T : IPackage
    {
        return Task.Run(loop);


        async void loop()
        {
            //TODO: 这里并发会有问题
            if (_singleReceive.ContainsKey(webSocket))
            {
                return;
            }
            else _singleReceive[webSocket] = true;


            while (webSocket.State == WebSocketState.Open)
            {
                try
                {
                    var raw_data = await receive_message(webSocket);

                    if (string.IsNullOrWhiteSpace(raw_data)) continue;

                    var package = JsonConvert.DeserializeObject<T>(raw_data) ??
                        throw new ArgumentException("数据包解析失败");

                    receivedCallback?.Invoke(package);
                }
                catch (Exception ex)
                {
                    if (excetionProcess?.Invoke(ex) != true) return;
                }
            }
        }

        static async Task<string> receive_message(WebSocket webSocket)
        {
            var receiveBuffer = new byte[1024];
            using MemoryStream dataStream = new();


            WebSocketReceiveResult result;
            do
            {
                result = await webSocket.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), CancellationToken.None);
                dataStream.Write(receiveBuffer, 0, result.Count);

            } while (!result.EndOfMessage);


            dataStream.Seek(0, SeekOrigin.Begin);

            return Encoding.UTF8.GetString(dataStream.ToArray());
        }
    }

    /// <summary>
    /// 发送一个包
    /// </summary>
    public static async Task SendAsync<T>(this WebSocket webSocket, T package) where T : IPackage
    {
        if (!_sendLock.ContainsKey(webSocket))
        {
            _sendLock[webSocket] = new(1, 1);
        }

        await _sendLock[webSocket].WaitAsync();

        try
        {
            var json = JsonConvert.SerializeObject(package);
            if (string.IsNullOrWhiteSpace(json)) return;

            ArraySegment<byte> data = new(Encoding.UTF8.GetBytes(json));
            await webSocket.SendAsync(data, WebSocketMessageType.Text, true, CancellationToken.None);
        }
        finally
        {
            _sendLock[webSocket].Release();
        }
    }
}
